import os
import torch

from typing import Dict, List, Any


class TensorIO:
    @staticmethod
    def save_tensors(tensors: Dict[str, List[Dict[tuple, torch.Tensor]]], cfg: Dict[str, Any]):
        save_path = cfg.get('input_tensor_output_folder')
        os.makedirs(save_path, exist_ok=True)
        
        for seq_id, frame_tensors in tensors.items():
            compressed_frames = []
            for tensor_dict in frame_tensors:
                if not tensor_dict:
                    break
                compressed_frames.append((
                    list(tensor_dict.keys()),
                    torch.stack(list(tensor_dict.values()))
                ))
            
            filename = f"{cfg.get('data_name')}_{seq_id}.pt"
            filepath = os.path.join(save_path, filename)
            
            torch.save(
                compressed_frames,
                filepath,
                _use_new_zipfile_serialization=True
            )

    @staticmethod
    def load_tensors(load_path: str) -> Dict[str, List[Dict[tuple, torch.Tensor]]]:
        tensors = {}
        
        files = sorted(
            [f for f in os.listdir(load_path) if f.endswith('.pt')],
            key=lambda x: int(x.split('_')[0])
        )
        
        for filename in files:
            filepath = os.path.join(load_path, filename)
            compressed_frames = torch.load(filepath)
            
            seq_id = '_'.join(filename.split('_')[1:]).replace('.pt', '')
        
            frame_tensors = []
            for keys, values in compressed_frames:
                frame_tensors.append({tuple(k): v for k, v in zip(keys, values)})
            
            tensors[seq_id] = frame_tensors
        
        return tensors
    
    @staticmethod
    def save_labels(labels: Dict[str, Dict[str, torch.Tensor]], cfg: Dict[str, Any]):
        save_path = cfg.get('label_input_tensor_output_folder')
        os.makedirs(save_path, exist_ok=True)
        
        for seq_id, class_tensors in labels.items():
            compressed_data = torch.stack(list(class_tensors.values()))
            
            filename = f"{cfg.get('data_name')}_{seq_id}.pt"
            filepath = os.path.join(save_path, filename)
            
            torch.save(
                compressed_data,
                filepath,
                _use_new_zipfile_serialization=True
            )
